'''
数据处理脚本 seg  windows linux
labelme格式转yoloseg格式
随机裁剪
'''
import datetime
import glob
import json
import os.path
import random
import shutil
import sys
import traceback

import cv2
import numpy as np


def celue_default(content, self):
    '''
    图像裁剪策略
    Args:
        t: one target
            {
              "label": "logo",
              "points": [
                [
                  702.0714285714287,
                  966.8571428571429
                ],
                ...
              ],
                }

    Returns:
        res_list4_listn:  [[int,int,int,int],,,] 裁剪框 xyxy
    '''
    Width, Height = content['imageWidth'], content['imageHeight']
    res_list4_listn = []

    # each target
    for ind, t in enumerate(content['shapes']):
        if self.ref.get(t['label']) is None:  # 不在self.ref里的标签不考虑
            continue
        # 1 获取裁剪框
        if t['label'] == 'bar' : # 策略3 20240511
            continue

        xys = np.array(t['points'])  # n*2
        x1, y1, x2, y2 = min(xys[:, 0]), min(xys[:, 1]), max(xys[:, 0]), max(xys[:, 1]),
        bw, bh = x2 - x1, y2 - y1

        scale = 5.0
        len_max = max(bw, bh)  # 目标大边
        # 随机裁
        x1_rand = random.randint(0, int(scale * len_max))
        x2_rand = random.randint(0, int(scale * len_max))
        y1_rand = random.randint(0, int(scale * len_max))
        y2_rand = random.randint(0, int(scale * len_max))
        res_list4_listn.append([int(max(x1 - x1_rand, 0)), int(max(y1 - y1_rand, 0)),
                                int(min(x2 + x2_rand, Width)), int(min(y2 + y2_rand, Height))])

        # # 策略2 20240510 固定裁
        # x1_rand = len_max / 8
        # x2_rand = len_max / 8
        # y1_rand = len_max / 8
        # y2_rand = len_max / 8
        # res_list4_listn.append([int(max(x1 - x1_rand, 0)), int(max(y1 - y1_rand, 0)),
        #                         int(min(x2 + x2_rand, Width)), int(min(y2 + y2_rand, Height))])

    # # 策略 特定roi裁剪
    # roi = [1410, 560, 2380, 1580]  # xyxy
    # res_list4_listn.append(roi)
    return res_list4_listn


def celue_trainV8Seg_strap_plate(content, self):
    '''
    图像裁剪策略
    Args:
        t: one target
            {
              "label": "logo",
              "points": [
                [
                  702.0714285714287,
                  966.8571428571429
                ],
                ...
              ],
                }

    Returns:
        res_list4_listn:  [[int,int,int,int],,,] 裁剪框 xyxy
    '''
    Width, Height = content['imageWidth'], content['imageHeight']
    res_list4_listn = []

    # each target
    for ind, t in enumerate(content['shapes']):
        if self.ref.get(t['label']) is None:  # 不在self.ref里的标签不考虑
            continue
        # 1 获取裁剪框
        # if t['label'] == 'bar' : # 策略3 20240511
        #     continue
        if not t['label'] == 'neicheng': # 只裁剪内衬框
            continue
        xys = np.array(t['points'])  # n*2
        x1, y1, x2, y2 = min(xys[:, 0]), min(xys[:, 1]), max(xys[:, 0]), max(xys[:, 1]),
        bw, bh = x2 - x1, y2 - y1

        scale = 1.0
        len_max = max(bw, bh)  # 目标大边
        # 随机裁
        repeat = 2
        for i in range(repeat):
            x1_rand = random.randint(0, int(scale * len_max))
            x2_rand = random.randint(0, int(scale * len_max))
            y1_rand = random.randint(0, int(scale * len_max))
            y2_rand = random.randint(0, int(scale * len_max))
            res_list4_listn.append([int(max(x1 - x1_rand, 0)), int(max(y1 - y1_rand, 0)),
                                    int(min(x2 + x2_rand, Width)), int(min(y2 + y2_rand, Height))])

        # # 策略2 20240510 固定裁
        # x1_rand = len_max / 8
        # x2_rand = len_max / 8
        # y1_rand = len_max / 8
        # y2_rand = len_max / 8
        # res_list4_listn.append([int(max(x1 - x1_rand, 0)), int(max(y1 - y1_rand, 0)),
        #                         int(min(x2 + x2_rand, Width)), int(min(y2 + y2_rand, Height))])

    # # 策略 特定roi裁剪
    # roi = [1410, 560, 2380, 1580]  # xyxy
    # res_list4_listn.append(roi)
    return res_list4_listn

def celue_phone(content, self): # per image
    '''
        计算每张图片的裁剪框
        Returns:
            res_list4_listn:  [[int,int,int,int],,,] 裁剪框 xyxy
    '''
    Width, Height = content['imageWidth'], content['imageHeight']
    res_list4_listn = []

    for ind, t in enumerate(content['shapes']):  # each target
        if self.ref.get(t['label']) is None:  # 不在self.ref里的标签不考虑
            continue

        # 1 获取裁剪框
        fix_scale = 1/8
        if t['label'] == 'logo': #
            fix_scale = 4

        xys = np.array(t['points'])  # n*2
        x1, y1, x2, y2 = min(xys[:, 0]), min(xys[:, 1]), max(xys[:, 0]), max(xys[:, 1]),
        bw, bh = x2 - x1, y2 - y1
        len_max = max(bw, bh)  # 目标大边


        # 策略2 20240510 固定裁, 外扩1/8
        x1_rand = len_max * fix_scale
        x2_rand = len_max * fix_scale
        y1_rand = len_max * fix_scale
        y2_rand = len_max * fix_scale

        x1_clip = int(max(x1 - x1_rand, 0))
        y1_clip = int(max(y1 - y1_rand, 0))
        x2_clip = int(min(x2 + x2_rand, Width))
        y2_clip = int(min(y2 + y2_rand, Height))
        res_list4_listn.append([x1_clip, y1_clip, x2_clip, y2_clip])

    # # 策略 特定roi裁剪
    # roi = [1410, 560, 2380, 1580]  # xyxy
    # res_list4_listn.append(roi)

    return res_list4_listn

def celue_trainV8Seg_fubahanxi(json_path, self): # per image
    '''

        Returns:
            res_list4_listn:  [[int,int,int,int],,,] 裁剪框 xyxy
    '''
    with open(json_path, 'r') as load_f:
        content = json.load(load_f)
    pre_dir_name = os.path.basename(os.path.dirname(json_path)) # 父文件夹名
    Width, Height = content['imageWidth'], content['imageHeight']
    img_name = content['imagePath']
    res_list4_listn = []

    for ind, t in enumerate(content['shapes']):  # each target
        if self.ref.get(t['label']) is None:  # 不在self.ref里的标签不考虑
            continue

        fix_scale = 2
        if t['label'] == 't0':  #
            continue

        xys = np.array(t['points'])  #  n*2
        x1, y1, x2, y2 = min(xys[:, 0]), min(xys[:, 1]), max(xys[:, 0]), max(xys[:, 1]),
        bw, bh = x2 - x1, y2 - y1
        len_max = max(bw, bh)  # 目标大边

        # 策略2 20240510 固定裁, 外扩1/8
        x1_rand = len_max * fix_scale
        x2_rand = len_max * fix_scale
        y1_rand = len_max * fix_scale
        y2_rand = len_max * fix_scale

        x1_clip = int(max(x1 - x1_rand, 0))
        y1_clip = int(max(y1 - y1_rand, 0))
        x2_clip = int(min(x2 + x2_rand, Width))
        y2_clip = int(min(y2 + y2_rand, Height))
        res_list4_listn.append([x1_clip, y1_clip, x2_clip, y2_clip])

    # 策略 特定roi裁剪
    if pre_dir_name in ("20240815", "20240815_cp2"):
        rois = [
            [1550, 1070, 1840, 1270],  # xyxy 013
            [2160, 1070, 2470, 1250],  # xyxy

            [4170, 1120, 4460, 1320],  # xyxy
            [4190, 1400, 4460, 1620],  # xyxy
            [4170, 1620, 4450, 1860],  # xyxy
            [4180, 1870, 4440, 2090],  # xyxy
        ]
        res_list4_listn.extend(rois)
    elif pre_dir_name in ('20241016hou_preroi','20240821_jdg','20240929','20241008'):
        rois = [
            [1240, 1990, 1530, 2230],
            [1220, 2220, 1490, 2470],
            [1230, 2480, 1490, 2700],
            [1240, 2740, 1510, 2940],
            [3280, 2780, 3540, 3000],
            [3880, 2800, 4180, 2980],
        ]
        res_list4_listn.extend(rois)

    elif pre_dir_name in ('20241016hou'):
        rois = [
            [1464, 1321, 1726, 1531],
            [1422, 1556, 1692, 1785],
            [1428, 1807, 1701, 2010],
            [1433, 2080, 1694, 2285],
            [3435, 2125, 3760, 2410],
            [3682, 1707, 4703, 2060],  # 5插槽
            [771, 1596, 1414, 1996],  # 2插槽
        ]
        res_list4_listn.extend(rois)


    else:
        rois = [
            [1480, 1400, 1750, 1600],
            [1460, 1600, 1750, 1850],
            [1470, 1860, 1730, 2060],
            [1480, 2150, 1770, 2350],
            [3500, 2210, 3790, 2420],

            [3670, 1670, 4750, 2160],
            [788, 1630, 1410, 2137],
        ]
        res_list4_listn.extend(rois)
    return res_list4_listn

def celue_trainV8Seg_hanQianXianShu(content, self): # per image
    '''
        焊前 线束
        Returns:
            res_list4_listn:  [[int,int,int,int],,,] 裁剪框 xyxy
    '''
    Width, Height = content['imageWidth'], content['imageHeight']
    img_name = content['imagePath']
    res_list4_listn = []

    for ind, t in enumerate(content['shapes']):  # each target
        if self.ref.get(t['label']) is None:  # 不在self.ref里的标签不考虑
            continue

        fix_scale = 2

        xys = np.array(t['points'])  #  n*2
        x1, y1, x2, y2 = min(xys[:, 0]), min(xys[:, 1]), max(xys[:, 0]), max(xys[:, 1]),
        bw, bh = x2 - x1, y2 - y1
        len_max = max(bw, bh)  # 目标大边

        # 策略2 20240510 固定裁, 外扩fix_scale
        x1_rand = len_max * fix_scale
        x2_rand = len_max * fix_scale
        y1_rand = len_max * fix_scale
        y2_rand = len_max * fix_scale

        x1_clip = int(max(x1 - x1_rand, 0))
        y1_clip = int(max(y1 - y1_rand, 0))
        x2_clip = int(min(x2 + x2_rand, Width))
        y2_clip = int(min(y2 + y2_rand, Height))
        res_list4_listn.append([x1_clip, y1_clip, x2_clip, y2_clip])

    # 策略 特定roi裁剪
    rois = [
        [2490, 2600, 3640, 3100],  # xyxy 线束

        [4145, 1899, 4424, 2137],  # xyxy 焊线焊盘
        [4191, 1664, 4467, 1875],  # xyxy
        [4162, 1429, 4445, 1642],  # xyxy
        [4151, 1169, 4448, 1394],  # xyxy
        [2175, 1067, 2462, 1299],  # xyxy

    ]
    res_list4_listn.extend(rois)

    return res_list4_listn

def celue_trainV8Seg_smallBxh(content, self): # per image
    '''
        小保险盒
        Returns:
            res_list4_listn:  [[int,int,int,int],,,] 裁剪框 xyxy
    '''
    Width, Height = content['imageWidth'], content['imageHeight']
    img_name = content['imagePath']
    res_list4_listn = []

    for ind, t in enumerate(content['shapes']):  # each target
        # if self.ref.get(t['label']) is not None:  # 在self.ref里的标签不考虑
        #     continue
        xys = np.array(t['points'])  # n*2
        x1, y1, x2, y2 = min(xys[:, 0]), min(xys[:, 1]), max(xys[:, 0]), max(xys[:, 1]),
        bw, bh = x2 - x1, y2 - y1
        len_max = max(bw, bh)  # 目标大边

        if t['label'] in ['roi_base','jiaxian']:
            for i in range(50):
                random_scale = 3.0
                # 随机裁
                x1_rand = random.randint(0, int(random_scale * len_max))
                x2_rand = random.randint(0, int(random_scale * len_max))
                y1_rand = random.randint(0, int(random_scale * len_max))
                y2_rand = random.randint(0, int(random_scale * len_max))
                res_list4_listn.append([int(max(x1 - x1_rand, 0)), int(max(y1 - y1_rand, 0)),
                                        int(min(x2 + x2_rand, Width)), int(min(y2 + y2_rand, Height))])

        else:
            if t['label'] not in ['box_muduanzi', 'box_pinzhen', 'box_baoxianpian', 'group_box_pinzhen',
                                  'roi_base']:  # 只考虑roi
                continue
            fix_scale = 0.2

            # 策略 20240510 固定裁, 外扩fix_scale
            x1_rand = len_max * fix_scale
            x2_rand = len_max * fix_scale
            y1_rand = len_max * fix_scale
            y2_rand = len_max * fix_scale

            x1_clip = int(max(x1 - x1_rand, 0))
            y1_clip = int(max(y1 - y1_rand, 0))
            x2_clip = int(min(x2 + x2_rand, Width))
            y2_clip = int(min(y2 + y2_rand, Height))
            res_list4_listn.append([x1_clip, y1_clip, x2_clip, y2_clip])
            res_list4_listn.append([x1_clip, y1_clip, x2_clip, y2_clip])

            # 策略 随机裁 1217
            for i in range(2):
                random_scale = (3.0,5.0)
                x1_rand = random.randint(int(random_scale[0] * len_max), int(random_scale[1] * len_max))
                x2_rand = random.randint(int(random_scale[0] * len_max), int(random_scale[1] * len_max))
                y1_rand = random.randint(int(random_scale[0] * len_max), int(random_scale[1] * len_max))
                y2_rand = random.randint(int(random_scale[0] * len_max), int(random_scale[1] * len_max))
                res_list4_listn.append([int(max(x1 - x1_rand, 0)), int(max(y1 - y1_rand, 0)),
                                        int(min(x2 + x2_rand, Width)), int(min(y2 + y2_rand, Height))])

    return res_list4_listn

def celue_trainV8Seg_lianlun(content, self):
    '''
    图像裁剪策略
    Args:
        t: one target
            {
              "label": "logo",
              "points": [
                [
                  702.0714285714287,
                  966.8571428571429
                ],
                ...
              ],
                }

    Returns:
        res_list4_listn:  [[int,int,int,int],,,] 裁剪框 xyxy
    '''
    Width, Height = content['imageWidth'], content['imageHeight']
    res_list4_listn = []

    # each target
    for ind, t in enumerate(content['shapes']):
        if self.ref.get(t['label']) is None:  # 不在self.ref里的标签不考虑
            continue
        # 1 获取裁剪框
        xys = np.array(t['points'])  # n*2
        x1, y1, x2, y2 = min(xys[:, 0]), min(xys[:, 1]), max(xys[:, 0]), max(xys[:, 1]),
        bw, bh = x2 - x1, y2 - y1

        scale = 10.0 # huaheng bengchi yuancao
        len_max = max(bw, bh)  # 目标大边
        # 随机裁
        repeat = 10
        for i in range(repeat):
            x1_rand = random.randint(0, int(scale * len_max))
            x2_rand = random.randint(0, int(scale * len_max))
            y1_rand = random.randint(0, int(scale * len_max))
            y2_rand = random.randint(0, int(scale * len_max))
            res_list4_listn.append([int(max(x1 - x1_rand, 0)), int(max(y1 - y1_rand, 0)),
                                    int(min(x2 + x2_rand, Width)), int(min(y2 + y2_rand, Height))])



    return res_list4_listn
class base:
    '''
    labelme格式 to yolo格式
    '''

    def __init__(self):
        pass

    def single_jsontotxt(self, json_path, out_path):
        '''
        seg: json转txt，单个文件，
        :param json_path:输入json文件地址
        :param out_path:输出txt文件地址  cls xywh
        :return:
        '''

        # load json
        try:
            with open(json_path, 'r') as load_f:
                content = json.load(load_f)
                # ret 2 seg
                for t in content['shapes']:
                    if t['shape_type'] == 'rectangle':
                        [x1, y1], [x2, y2] = t['points']
                        t['points'] = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]
                        t['shape_type'] = 'polygon'
        except:
            print(json_path)
            print(traceback.format_exc())
            raise

        # each box
        file_str = ''
        W, H = content['imageWidth'], content['imageHeight']
        for t in content['shapes']:

            # type, (标签名转typeid) 修改它
            # ref = {'yuanzhu_hong': 0, 'yuanzhu_lan': 1, 'fangkuai_hong': 2}
            if self.ref.get(t['label']) is None: # 不在self.ref里的标签不考虑
                continue
            type = self.ref[t['label']]
            file_str += str(type)


            for x, y in t['points']:
                x = x / W
                y = y / H
                file_str += ' ' + str(round(x, 6)) + ' ' + str(round(y, 6))
            file_str += '\n'

        # save
        filename = out_path
        if os.path.exists(filename):
            os.remove(filename)
        try:
            os.mknod(filename)  # win不支持
            fp = open(filename, mode="r+", encoding="utf-8")
        except:
            # print('非 linux平台')
            fp = open(filename, mode="w", encoding="utf-8") # win
        fp.write(file_str[:-1])
        fp.close()

    @staticmethod
    def txt2json(txt_path,json_path):
        '''
        分割 yolo格式转labelme格式
        Args:
            txt_path:
            json_path:

        Returns:

        '''
        imagePath = txt_path[:-3]+'jpg'
        img = cv2.imread(imagePath)
        content_json = {
            "version": "5.4.1",
            "flags": {},
            "shapes": [],
            "imagePath": imagePath,
            "imageData": None,
            "imageHeight": img.shape[0],
            "imageWidth": img.shape[1]
        }
        # 打开txt
        with open(txt_path, 'r', encoding='utf-8') as file:
            content_txt = file.read()
            print(content_txt)
        its = content_txt.split('\n')
        for it in its:
            values = it.split(' ')
            shape = {
                "label": values[0],
                "points": [
                ],
                "group_id": None,
                "description": "",
                "shape_type": "polygon",
                "flags": {},
                "mask": None
            }
            numb = (len(values)-1)//2
            for i in range(numb):
                x, y = float(values[2*i+1]), float(values[2*i+2])
                shape['points'].append([x*img.shape[1], y*img.shape[0]])
            content_json['shapes'].append(shape)

        with open(json_path, 'w', encoding='utf-8') as file:
            json.dump(content_json, file, indent=4,  ensure_ascii=False)



    def do_augment(self, json_path, aug_name='', isClip=True, scale = 1.0):
        # load json
        with open(json_path, 'r') as load_f:
            content = json.load(load_f)
            # box 2 seg
            for st in content['shapes']:
                if st['shape_type'] == 'rectangle':
                    [x1, y1], [x2, y2] = st['points']
                    st['points'] = [x1, y1], [x2, y1], [x2, y2], [x1, y2]
                    st['shape_type'] = 'polygon'
        W, H = content['imageWidth'], content['imageHeight']

        if isClip:
            try:
                # 1 获取裁剪框
                if aug_name == 'trainV8Seg_phone':
                    res_list4_listn = celue_phone(content, self)  #
                    # res_list4_listn = celue_phone(t, Width=W, Height=H) #
                elif aug_name == 'trainV8Seg_fubahanxi':
                    res_list4_listn = celue_trainV8Seg_fubahanxi(json_path, self)
                elif aug_name == 'trainV8Seg_hanQianXianShu':
                    res_list4_listn = celue_trainV8Seg_hanQianXianShu(content, self)
                elif aug_name == 'trainV8Seg_smallBxh':
                    res_list4_listn = celue_trainV8Seg_smallBxh(content, self)
                elif aug_name == 'trainV8Seg_strap_plate':
                    res_list4_listn = celue_trainV8Seg_strap_plate(content, self)
                elif aug_name == 'trainV8Seg_lianlun':
                    res_list4_listn = celue_trainV8Seg_lianlun(content, self)# celue_trainV8Seg_lianlun
                else:
                    res_list4_listn = celue_default(content, self)

                # 2 裁剪图片和标签
                for clip_ind, (x1_clip,y1_clip,x2_clip,y2_clip) in enumerate(res_list4_listn): # per 策略

                    img = cv2.imread(json_path[:-5] + '.jpg')
                    img_clip = img[y1_clip:y2_clip, x1_clip:x2_clip]  # 保存裁剪图片

                    lab_str = ''
                    for st in content['shapes']: # 统计裁切框内的所有标签
                        # type = self.ref[t['label']]
                        if self.ref.get(st['label']) is None:  # 不在self.ref里的标签不考虑
                            continue
                        type = self.ref[st['label']]  # bug修正 20240305

                        # # box 2 seg
                        # if st['shape_type'] == 'rectangle':
                        #     [x1, y1], [x2, y2] = st['points']
                        #     st['points'] = [x1, y1], [x2, y1], [x2, y2], [x1, y2]
                        # else:
                        #     st['points'] = st['points']

                        # 判断目标是否在img_clip外
                        cnt = 0
                        for x, y in st['points']:
                            x = (x - x1_clip) / (x2_clip - x1_clip)
                            y = (y - y1_clip) / (y2_clip - y1_clip)
                            if x < 0 or x > 1 or y < 0 or y > 1:  # 20250125 在clip图片外
                                cnt += 1
                        # if cnt > len(st['points'])*0.5 : #如果目标有一半的点在clip图片外，不考虑该目标
                        #     continue
                        if cnt > len(st['points'])*0.0: #如果目标有点在clip图片外，不考虑该目标
                            continue
                        lab_str += str(type)  # 记录标签，
                        for x, y in st['points']:
                            x = (x - x1_clip) / (x2_clip - x1_clip)
                            y = (y - y1_clip) / (y2_clip - y1_clip)
                            x = min(max(x,0),1)
                            y = min(max(y,0),1) # 0-1

                            lab_str += ' ' + str(round(x, 6)) + ' ' + str(round(y, 6))
                        lab_str += '\n'
                    # 3 保存
                    # js_name = json_path.split('\\')[-1]  # wind
                    # save_dir = json_path[:-len(js_name) - 1] + '_mini'  # window
                    # wind or linux
                    js_name = os.path.basename(json_path)
                    save_dir = os.path.dirname(json_path) + '_mini'
                    if save_dir[-8:] == 'ini_mini':
                        pass
                    if not os.path.exists(save_dir):
                        os.mkdir(save_dir)

                    ijpg_path = f'{save_dir}/{js_name[:-5]}_clip_{clip_ind}.jpg'
                    itxt_path = f'{save_dir}/{js_name[:-5]}_clip_{clip_ind}.txt'
                    ijson_path = f'{save_dir}/{js_name[:-5]}_clip_{clip_ind}.json'
                    cv2.imwrite(ijpg_path, img_clip)
                    with open(itxt_path, mode="w", encoding="utf-8") as fp: # win
                        fp.write(lab_str[:-1])
                    # self.txt2json(txt_path, json_path) # 转为json, 用于查看

            except Exception as e:
                print(e)
                print(json_path)
                print(traceback.format_exc())
                raise Exception('do_augment 错误')

    def make_txt_from_json(self):
        '''
        json to txt
        '''
        print(f'Make_txt_from_json...')
        ls = glob.glob(self.json)
        total = 0
        for i in ls:
            pre_name = os.path.basename(os.path.dirname(i))
            if pre_name[-4:] == 'mini':
                continue
            txt_path = f'{i[:-5]}.txt'
            self.single_jsontotxt(i, txt_path)  # 转换
            # self.do_augment(i) # 随机裁剪
            total += 1
            # print(f'{total}/{int(len(ls))}: saved to {txt_path}')
        print(f'Make_txt_from_json , Done.')

    def augment_from_json(self, aug_name='',random_scale = 1.0): # 随机裁剪
        '''
        随机裁剪,生成 ~_mini文件夹
        '''
        print(f'Augment: random clip...')
        ls = glob.glob(self.json)
        total = 0
        for i in ls:
            pre_name = os.path.basename(os.path.dirname(i))
            print(pre_name)
            if pre_name[-4:] == 'mini':
                continue
            self.do_augment(i, aug_name=aug_name, scale=random_scale) # 随机裁剪
            total += 1
            # print(f'{total}/{int(len(ls))}: saved to {txt_path}')
        print(f'Augment: random clip, Done.')
    # def make_yolov5(self):
    #     pass

    def update_yolov5(self):
        '''
        make yolov5 or add a unit to yolov5
        '''

        # source
        self.info()

        # dst
        if not os.path.exists(self.yolov5):
            os.makedirs(self.yolov5)
        else:
            print(f'rm {self.yolov5}/*/*...')
            if os.path.exists(f'{self.yolov5}images'):
                shutil.rmtree(f'{self.yolov5}images') # 删除，重新转换,目前不能增量式添加数据
            if os.path.exists(f'{self.yolov5}labels'):
                shutil.rmtree(f'{self.yolov5}labels')
        out_fir_tra_img = self.yolov5 + 'images/train/'
        out_fir_val_img = self.yolov5 + 'images/val/'
        out_fir_tra_txt = self.yolov5 + 'labels/train/'
        out_fir_val_txt = self.yolov5 + 'labels/val/'
        iltv_dir = [out_fir_tra_img, out_fir_val_img, out_fir_tra_txt, out_fir_val_txt]
        for i in iltv_dir:
            if not os.path.exists(i): # 如果不存在
                os.makedirs(i)

        print('before:')
        t_i, v_i, t_tx, v_tx = len(os.listdir(out_fir_tra_img)), len(os.listdir(out_fir_val_img)), \
                               len(os.listdir(out_fir_tra_txt)), len(os.listdir(out_fir_val_txt))
        print(f'{out_fir_tra_img}:{t_i}')
        print(f'{out_fir_val_img}:{v_i}')
        print(f'{out_fir_tra_txt}:{t_tx}')
        print(f'{out_fir_val_txt}:{v_tx}')

        # 划分 各类随机抽1%
        ls_img = glob.glob(self.jpg)
        random.shuffle(ls_img)
        flag = max(int(0.2 * len(ls_img)) + 1, 1)  # 20%
        # ls_img_val = ls_img[:flag]
        # ls_img_tra = ls_img[flag:]
        ls_img_val = ls_img[:flag]
        ls_img_tra = ls_img[0:]
        # ls_img_tra = ls_img[flag:]

        # 添加至验证集

        for i in ls_img_val:
            # jpg
            # nam = i.split('/')[-1]
            # nam = i.split('\\')[-1] # windows
            # self.buff = i.split('\\')[-2] # user
            nam = os.path.basename(i) # wind, linux
            self.buff = os.path.basename(os.path.dirname(i)) # wind, linux
            shutil.copy(i, out_fir_val_img + f'{self.buff}_{nam}')  # 添加至val_img
            # lab
            # txt_abpath = self.txt + nam[:-4] + '.txt'  #
            txt_abpath = f'{i[:-4]}.txt'  #
            if os.path.exists(txt_abpath):  # 如果有标签文件
                shutil.copy(txt_abpath, out_fir_val_txt + f'{self.buff}_{nam[:-4]}.txt')

        # 添加至训练集
        for ind, i in enumerate(ls_img_tra):
            # jpg
            # nam = i.split('/')[-1]
            # nam = i.split('\\')[-1]  # windows
            # self.buff = i.split('\\')[-2]  # user 父文件名为前缀
            nam = os.path.basename(i) # wind, linux
            self.buff = os.path.basename(os.path.dirname(i)) # wind, linux
            shutil.copy(i, out_fir_tra_img + f'{self.buff}_{nam}')  # 添加至tra_img
            # lab
            txt_abpath = f'{i[:-4]}.txt'
            if os.path.exists(txt_abpath):
                shutil.copy(txt_abpath, out_fir_tra_txt + f'{self.buff}_{nam[:-4]}.txt')  # 添加至tra_lab

        # check
        print('after:')
        t_i_d, v_i_d, t_tx_d, v_tx_d = len(os.listdir(out_fir_tra_img)), len(os.listdir(out_fir_val_img)), \
                                       len(os.listdir(out_fir_tra_txt)), len(os.listdir(out_fir_val_txt))
        print(f'{out_fir_tra_img}:{t_i_d}')
        print(f'{out_fir_val_img}:{v_i_d}')
        print(f'img added: {t_i_d + v_i_d - t_i - v_i}')
        print(f'{out_fir_tra_txt}:{t_tx_d}')
        print(f'{out_fir_val_txt}:{v_tx_d}')
        print(f'lab added: {t_tx_d + v_tx_d - t_tx - v_tx}')

        print(f'Update_yolov5, Done.')

    def remove_yolov5(self, buff):
        '''
        remove a unit from yolov5
        :return:
        '''
        ls = glob.glob(f'{self.yolov5}*/*/*')
        for ind, i in enumerate(ls):
            # if i.split('/')[-1][:len(buff)] == buff:  # 根据前缀删除
            name = os.path.basename(i) # wind, linux
            if name[:len(buff)] == buff:  # windows
                os.remove(i)
                print(f'{ind}/{len(ls)}:{i}')

    def info(self):
        # ls = glob.glob(self.glob_str)
        # print(self.buff)
        # print(
        #     f'origin:{len(ls)}\njpg:{len(os.listdir(self.jpg))}\njson:{len(os.listdir(self.json))}\ntxt:{len(os.listdir(self.txt))}\n')
        pass

    def info_txt(self):
        pass

    def resize_imgs(self, target_size = 416): # 缩小训练集图片分辨率，加快训练

        imgs_glob = self.yolov5+'images/*/*.jpg'
        ls = glob.glob(imgs_glob)
        for i in ls:
            img = cv2.imread(i)
            scale = target_size/max(img.shape[:2])
            # print(scale)
            try:
                img = cv2.resize(img, dsize=None,fx=scale,fy=scale)
            except:
                print(i)
            cv2.imwrite(i, img)
        print(f'resized done, total {len(ls)}, scale {scale}')



# class labelme2yoloseg(base):
#     jpg = r'D:\data\231207huoni\trainseg_zhihe\add_imgs\NOREAD_20231230_150531_407\*.jpg'
#     json = r'D:\data\231207huoni\trainseg_zhihe\add_imgs\NOREAD_20231230_150531_407\*.json'
#     txt = r'D:\data\231207huoni\trainseg_zhihe\add_imgs\NOREAD_20231230_150531_407\*.txt'
#     yolov5 = r'D:\data\231207huoni\trainseg_zhihe\imglabs_yolov5/'  # yolov5数据集地址
#     buff = 'd2024_0102'
#     # ref = {'zhihe': 0, 'zhuji_zheng': 1, 'zhuji_fan': 2, 'dianchi_zheng': 3, 'dianchi_fan': 4, 'xianshu': 5, 'shuomingshu': 6, 'wandai': 7 }
#     ref = {'zhihe': 0, 'logo': 1 }
#
#
#     def run(self):
#         self.make_txt_from_json()
#         self.update_yolov5()
#         # self.info()
#         self.resize_imgs(scale=0.3) # (2048, 3072, 3)


class labelme2yoloseg(base):
    def __init__(self):
        self.jpg = r'D:\data\231207huoni\trainseg_zhihe\add_imgs\NOREAD_20231230_150531_407\*.jpg'
        self.json = r'D:\data\231207huoni\trainseg_zhihe\add_imgs\NOREAD_20231230_150531_407\*.json'
        self.txt = r'D:\data\231207huoni\trainseg_zhihe\add_imgs\NOREAD_20231230_150531_407\*.txt'
        self.yolov5 = r'D:\data\231207huoni\trainseg_zhihe\imglabs_yolov5/'  # yolov5数据集地址
        self.buff = 'd2024_0102'
        # ref = {'zhihe': 0, 'zhuji_zheng': 1, 'zhuji_fan': 2, 'dianchi_zheng': 3, 'dianchi_fan': 4, 'xianshu': 5, 'shuomingshu': 6, 'wandai': 7 }
        self.ref = {'zhihe': 0, 'logo': 1}
        self.scale = 1.0

    def run(self, isAug = False, aug_name='', random_scale = 1.0):
        self.make_txt_from_json()
        if isAug:
            self.augment_from_json(aug_name=aug_name, random_scale=random_scale) # 随机裁剪
        self.update_yolov5()
        # self.info()
        self.resize_imgs(target_size=self.target_size)  # (2048, 3072, 3)

if __name__ == '__main__':
    pass
    # base().txt2json(r"D:\data\240815fubahanxi\trainV8Seg_hanQianXianShu\add_imgs\20241014_mini\1009-000156_clip_0.txt",
    #                 r"D:\data\240815fubahanxi\trainV8Seg_hanQianXianShu\add_imgs\20241014_mini\1009-000156_clip_0.json")